import copy
import itertools
from matplotlib import pyplot as plt
from matplotlib.patches import Rectangle
from matplotlib import animation

T = [[2, 3], [5, 4], [9, 6], [4, 7], [8, 1], [7, 2]]


def draw_point(data):
    x, y = [], []
    for p in data:
        x.append(p[0])
        y.append(p[1])
    plt.plot(x, y, 'bo')


def draw_line(xy_list):
    for xy in xy_list:
        x, y = xy
        plt.plot(x, y, 'g', lw=2)


def draw_square(square_list):
    current_axis = plt.gca()
    colors = itertools.cycle(['r', 'b', 'g', 'm', 'y', "#EB70AA", "#0099FF"])
    for square in square_list:
        current_axis.add_patch(
            Rectangle((square[0][0], square[1][0]), square[1][0] - square[0][1], square[1][1] - square[0][1],
                      color=next(colors)))


def median(data):
    m = int(len(data) / 2)
    return data[m], m


history_square = []


def build_kdtree(data, d, square):
    history_square.append(square)
    data = sorted(data, key=lambda x: x[d])
    p, m = median(data)
    del data[m]
    if m >= 0:
        sub_square = copy.deepcopy(square)
        if d == 0:
            sub_square[1][0] = p[0]
        else:
            sub_square[1][1] = p[1]
        history_square.append(sub_square)
        if m > 0:
            build_kdtree(data[:m], not d, sub_square)
    if len(data) > 1:
        sub_square = copy.deepcopy(square)
        if d == 0:
            sub_square[0][0] = p[0]
        else:
            sub_square[0][1] = p[1]

        build_kdtree(data[m:], not d, sub_square)


build_kdtree(T, 0, [[0, 0], [10, 10]])

fig = plt.figure()
aux = plt.axes(xlim=(0, 2), ylim=(-2, -2))
line, = aux.plot([], [], 'g', lw=2)
label = aux.text([], [], '')


def init():
    plt.axis([0, 10, 0, 10])
    plt.grid(True)
    plt.xlabel("x_1")
    plt.ylabel('x_2')
    plt.title("build kd tree")
    draw_point(T)


current_axis = plt.gca()
colors = itertools.cycle(["#FF6633", "g", "#3366FF", "c", "m", "y", '#EB70AA', '#0099FF', '#66FFFF'])


def animate(i):
    square = history_square[i]
    current_axis.add_patch(
        Rectangle((square[0][0], square[1][0]), square[1][0] - square[0][1], square[1][1] - square[0][1],
                  color=next(colors)))
    return


#anim = animation.FuncAnimation(fig, animate, init_func=init, frames=len(history_square), interval=1000, blit=False)
init()
plt.show()
#anim.save('kdtree_build.gif', fps=2, writer="imagemagick")
